# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
CV2 图片画矩形框、直线、写字。。。

# 创建时间: 2019/7/8 20:58
# 开发作者: vincent
# 创建平台: PyCharm Community Edition
# 版    本: V1.0
"""
import cv2
import numpy as np


# 在图片上画矩形框及打印字符串
def draw_rectangle(image, save_name, text_string=None, text_xy=None, text_font=None, text_size=None,
                   text_thickness=None, text_color=(255, 0, 0), rectangle_xy=None, rectangle_fill=None,
                   rectangle_color=(255, 0, 0), rectangle_width=None):
    """
    :param image:            numpy 格式  (1000, 2000, 3)   通道序列：BGR
    :param save_name:        图片保存路径
    :param text_string:      打印字符内容 list/tuple  ["test"],("test")
    :param text_xy:          文本打印起始位置  None时默认打印在上方或下方
    :param text_font:        中文所需的字体库  "./platech.ttf"， None时使用CV画图无法显示中文
    :param text_size:        文本大小
    :param text_thickness:   字符笔划大小 CV2 使用
    :param text_color:       文本显示颜色
    :param rectangle_xy:     矩形框位置  [(x1, y1, x2, y2),(x1, y1, x2, y2)]
    :param rectangle_fill:   矩形框填充颜色
    :param rectangle_color:  矩形框颜色
    :param rectangle_width:  矩形框线条大小
    :return:
    """
    if rectangle_width is None:
        rectangle_width = int(image.shape[1]/500)

    # PIL 画框、显示字符（支持中文, 需要指定中文字库路径 ）  输入图片通道为 RGB
    if text_font is not None:
        import cv2
        from PIL import Image, ImageDraw, ImageFont

        pil_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # RGB 通道转换
        pil_image = Image.fromarray(pil_image)

        # # PIL图片上打印汉字
        draw = ImageDraw.Draw(pil_image)

        if rectangle_xy is not None:
            # 边框颜色
            if not isinstance(rectangle_color[0], (list, tuple)):
                rectangle_color = [rectangle_color] * len(rectangle_xy)
            else:
                assert len(rectangle_color) == len(rectangle_xy), \
                    "Error: rectangle_color {} must be same length  with rectangle_xy or like (0,0,0) ".format(
                        rectangle_color)

            # 边框大小
            if not isinstance(rectangle_width, (list, tuple)):
                rectangle_width = [rectangle_width] * len(rectangle_xy)
            else:
                assert len(rectangle_width) == len(rectangle_xy), \
                    "Error: rectangle_width {} must be same length  with rectangle_xy or like 10 ".format(
                        rectangle_width)

            # 循环画每个矩形框
            for rectangle_xy_i, color_i, width_i in zip(rectangle_xy, rectangle_color, rectangle_width):
                draw.rectangle(rectangle_xy_i, rectangle_fill, color_i, width_i)

        if text_string is not None:
            if not isinstance(text_string, (list, tuple)):
                text_string = [text_string]

            if text_size is None:
                text_size = int(image.shape[1] / 20)

            # 若无指定文字位置，默认都在最上方
            if text_xy is None:
                text_xy = [(0, 0)] * len(text_string)
            elif not isinstance(text_xy[0], (list, tuple)):
                text_xy = [text_xy] * len(text_string)
            else:
                assert len(text_xy) == len(text_string), \
                    "Error: text_xy {} must be same length  with text_string or like (0,0) ".format(text_xy)

            # 文字大小
            if not isinstance(text_size, (list, tuple)):
                text_size = [text_size] * len(text_string)
            else:
                assert len(text_size) == len(text_string), \
                    "Error: text_size {} must be same length  with text_string or like 10 ".format(text_size)

            # 字体颜色
            if not isinstance(text_color[0], (list, tuple)):
                text_color = [text_color] * len(text_string)
            else:
                assert len(text_color) == len(text_string), \
                    "Error: text_color {} must be same length  with text_string or like (0,0,0) ".format(text_color)

            # 循环标注各个字符串
            for text_i, text_xy_i, text_size_i, text_color_i in zip(text_string, text_xy, text_size, text_color):
                font = ImageFont.truetype(font=text_font, size=text_size_i, encoding="utf-8")
                draw.text(text_xy_i, text_i, text_color_i, font=font)

        # 保存图片
        pil_image.save(save_name)
        return True

    # CV2 画框、显示字符（不支持中文）  输入图片通道为 BGR
    else:
        import copy
        import cv2

        # img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
        cv_image = copy.copy(image)

        if rectangle_xy is not None:
            # 边框颜色
            if not isinstance(rectangle_color[0], (list, tuple)):
                rectangle_color = [rectangle_color] * len(rectangle_xy)
            else:
                assert len(rectangle_color) == len(rectangle_xy), \
                    "Error: rectangle_color {} must be same length  with rectangle_xy or like (0,0,0) ".format(
                        text_color)

            # 边框大小
            if not isinstance(rectangle_width, (list, tuple)):
                rectangle_width = [rectangle_width] * len(rectangle_xy)
            else:
                assert len(rectangle_width) == len(rectangle_xy), \
                    "Error: rectangle_width {} must be same length  with rectangle_xy or like 10 ".format(
                        rectangle_width)

            # 循环画每个矩形框
            for rectangle_xy_i, color_i, width_i in zip(rectangle_xy, rectangle_color, rectangle_width):
                cv2.rectangle(cv_image, rectangle_xy_i[0:2], rectangle_xy_i[2:4], color_i, width_i)

        if text_string is not None:
            if not isinstance(text_string, (list, tuple)):
                text_string = [text_string]

            if text_size is None:
                text_size = image.shape[1] / 200

            if text_thickness is None:
                text_thickness = int(image.shape[1] / 200)

            # 若无指定文字位置，默认都在最下方
            if text_xy is None:
                text_xy = [(0, image.shape[0])] * len(text_string)
            elif not isinstance(text_xy[0], (list, tuple)):
                text_xy = [text_xy] * len(text_string)
            else:
                assert len(text_xy) == len(text_string), \
                    "Error: text_xy {} must be same length  with text_string or like (0,0) ".format(text_xy)

            # 文字大小
            if not isinstance(text_size, (list, tuple)):
                text_size = [text_size] * len(text_string)
            else:
                assert len(text_size) == len(text_string), \
                    "Error: text_size {} must be same length  with text_string or like 10 ".format(text_size)

            # 笔划粗细
            if not isinstance(text_thickness, (list, tuple)):
                text_thickness = [text_thickness] * len(text_string)
            else:
                assert len(text_thickness) == len(text_string), \
                    "Error: text_thickness {} must be same length  with text_string or like 10 ".format(text_thickness)

            # 字体颜色
            if not isinstance(text_color[0], (list, tuple)):
                text_color = [text_color] * len(text_string)
            else:
                assert len(text_color) == len(text_string), \
                    "Error: text_color {} must be same length  with text_string or like (0,0,0) ".format(text_color)

            # 循环标注各个字符串
            for text_i, text_xy_i, text_size_i, text_color_i, text_thickness_i in zip(text_string, text_xy, text_size,
                                                                                      text_color, text_thickness):

                cv2.putText(cv_image, text_i, text_xy_i, cv2.FONT_HERSHEY_PLAIN, text_size_i, text_color_i,
                            text_thickness_i)

        cv2.imencode(".jpg", cv_image)[1].tofile(save_name)
        return cv_image


# 多边形填充框
class FillPoly(object):
    def __init__(self, poly_points, img_size, color=(0, 100, 0)):
        """
        :param poly_points:  多边形各个顶点
        :param img_size:   图片大小
        :param color:   填充颜色
        """
        # 多边形填充
        self.base_poly_img = cv2.fillPoly(np.zeros(img_size), poly_points, color)

        # 多边形各顶点
        self.poly_points = poly_points

        # 模板图片大小
        self.img_size = img_size

    def __call__(self, img, line_color=None):
        """
        :param img:    待处理图片
        :param line_color:    边框颜色
        :return:
        """
        # 判断图片大小是否一致
        assert img.shape == self.img_size, \
            print('ImageShapeError:  BaseImg：{} InputImg：{}'.format(self.img_size, img.shape))
        # img = np.where(self.base_poly_img > 0, self.base_poly_img, img)
        img = self.base_poly_img + img
        img = np.where(img > 255, 255, img)

        # 判断是否画多边形边框
        if line_color is not None:
            img = cv2.polylines(img, self.poly_points, 1, line_color)

        return img


def main():
    # cv2读取图片
    img_path = 'test.jpg'
    img = cv2.imdecode(np.fromfile(img_path, dtype='uint8'), -1)
    print(img.shape)

    rectangle_xy = [(10, 10, 100, 200), (100, 100, 300, 330)]
    rectangle_color = (255, 0, 0)
    rectangle_width = [5, 1]
    text_xy = [(300, 300), (100, 100)]
    text_size = 5
    text_thickness = [5, 3]
    text_string = [u"您好 1", "Hello 2"]
    text_color = (255, 0, 0)

    # 画矩形框及文字
    img_new = draw_rectangle(img, "_test0.jpg", text_font='platech.ttf', text_xy=text_xy, text_string=text_string,
                             rectangle_xy=rectangle_xy, rectangle_color=rectangle_color,
                             rectangle_width=rectangle_width, text_size=30, text_color=text_color,
                             text_thickness=text_thickness)

    # 画矩形框及文字
    img_new = draw_rectangle(img, "_test1.jpg", text_font=None, text_xy=text_xy, text_string=text_string,
                             rectangle_xy=rectangle_xy, rectangle_color=rectangle_color,
                             rectangle_width=rectangle_width, text_size=text_size, text_color=text_color,
                             text_thickness=text_thickness)

    # 填充多边形
    points = np.array([[[100, 100], [500, 100], [200, 230], [150, 200], [100, 220]]], dtype='int32')
    fill_ploy = FillPoly(points, img_new.shape)
    cv2.imencode(".jpg", fill_ploy(img_new, line_color=(0, 0, 0)))[1].tofile('_test2.jpg')

if __name__ == '__main__':
    main()
